package com.yl.flink.sink;

import com.yl.util.FlinkUtil;
import com.yl.util.InfluxUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.flink.api.java.utils.ParameterTool;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.sink.RichSinkFunction;
import org.apache.flink.streaming.api.functions.sink.SinkFunction;
import org.apache.flink.streaming.connectors.influxdb.InfluxDBConfig;
import org.apache.flink.streaming.connectors.influxdb.InfluxDBPoint;
import org.apache.flink.util.CollectionUtil;
import org.apache.flink.util.StringUtils;
import org.influxdb.InfluxDB;
import org.influxdb.dto.Point;

import java.util.concurrent.TimeUnit;


/**
 * 重写 influxdb-connector，自定义 TimeUnit
 */
@Slf4j
public class CustomInfluxSink extends RichSinkFunction<InfluxDBPoint> {
    private transient InfluxDB influxDBClient;
    private final InfluxDBConfig influxDBConfig;
    private final TimeUnit timeUnit;

    public CustomInfluxSink(ParameterTool params) {
        this.influxDBConfig = FlinkUtil.getInfluxDBConfig(params);
        this.timeUnit = TimeUnit.MICROSECONDS;
    }

    @Override
    public void open(Configuration parameters) throws Exception {
        super.open(parameters);
        this.influxDBClient = InfluxUtil.connect(this.influxDBConfig);
    }

    @Override
    public void invoke(InfluxDBPoint dataPoint, SinkFunction.Context context) throws Exception {
        if (StringUtils.isNullOrWhitespaceOnly(dataPoint.getMeasurement())) {
            throw new RuntimeException("No measurement defined");
        } else {
            Point.Builder builder = Point
                    .measurement(dataPoint.getMeasurement())
                    .time(dataPoint.getTimestamp(), this.timeUnit);
            if (!CollectionUtil.isNullOrEmpty(dataPoint.getFields())) {
                builder.fields(dataPoint.getFields());
            }
            if (!CollectionUtil.isNullOrEmpty(dataPoint.getTags())) {
                builder.tag(dataPoint.getTags());
            }
            Point point = builder.build();
            this.influxDBClient.write(point);
        }
    }

    @Override
    public void close() {
        if (this.influxDBClient.isBatchEnabled()) {
            this.influxDBClient.disableBatch();
        }
        this.influxDBClient.close();
    }
}
